import torch.nn as nn
import torch
from PIL import Image
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter

from torch练习.torchvision中的dataset.nn_maxpool import original_img


class MyRelu(nn.Module):
    def __init__(self):
        super(MyRelu, self).__init__()
        # inplace是否对原数据进行替换
        # self.relu = nn.ReLU(inplace=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        return self.sigmoid(x)

# input = torch.tensor([
#     [1, -0.5],
#     [-1, 3]
# ])
# input = input.reshape((-1, 1, 2, 2))
# myrelu = MyRelu()
# output = myrelu(input)
# print(output)

myrelu = MyRelu()

original_img = Image.open("./dataset/mashu.jpg")
tensor_img = transforms.ToTensor()(original_img)
tensor_img = tensor_img.reshape(-1, 1, 1024, 1024)
output_img = myrelu(tensor_img)

writer = SummaryWriter("./logs")
writer.add_image("sigmoid 麻薯", output_img[1])
writer.close()